use jni::objects::{JClass, JObject, JString, JValue};
use jni::sys::jlong;
use jni::JNIEnv;
use std::sync::Arc;

/// LLModel trait - 用于处理大模型的流式输出
/// Java 侧需要实现这个接口
pub trait LLModel {
    /// 获取下一个模型输出，返回 None 表示流结束
    fn next(&mut self) -> Option<String>;
}

/// AgentResponseStream trait - 用于将 agent 结果返回到前端
/// Java 侧需要实现这个接口
pub trait AgentResponseStream {
    /// 更新流式输出
    fn update(&self, content: &str);
}

/// QAAgent 结构体
pub struct QAAgent {
    workspace: String,
}

impl QAAgent {
    /// 创建新的 QAAgent 实例
    pub fn new(workspace: String) -> Self {
        QAAgent { workspace }
    }

    /// 处理用户查询
    /// 
    /// # Arguments
    /// * `query` - 用户的输入
    /// * `model` - 大模型接口，用于获取流式输出
    /// * `stream` - 响应流，用于将结果返回到前端
    pub fn answer(
        &self,
        query: String,
        model: &mut dyn LLModel,
        stream: &dyn AgentResponseStream,
    ) {
        stream.update(&format!("Workspace: {}\n", self.workspace));
        stream.update(&format!("Query: {}\n", query));
        stream.update("Processing...\n");

        // 从模型获取流式输出
        while let Some(output) = model.next() {
            stream.update(&format!("Model output: {}\n", output));
        }

        stream.update("Done.\n");
    }
}

// ========== JNI 绑定部分 ==========

use jni::objects::GlobalRef;

/// JNI 包装结构，用于在 JNI 层存储 QAAgent
struct QAAgentWrapper {
    agent: Arc<QAAgent>,
}

/// Java LLModel 的 JNI 包装
/// 注意：这个结构持有 Java 对象的全局引用
struct JavaLLModel {
    obj: GlobalRef,
}

unsafe impl Send for JavaLLModel {}
unsafe impl Sync for JavaLLModel {}

impl JavaLLModel {
    fn new(env: &mut JNIEnv, obj: JObject) -> Result<Self, jni::errors::Error> {
        let global_ref = env.new_global_ref(obj)?;
        Ok(JavaLLModel { obj: global_ref })
    }

    fn next(&mut self, env: &mut JNIEnv) -> Option<String> {
        let result = env.call_method(
            self.obj.as_obj(),
            "next",
            "()Ljava/lang/String;",
            &[],
        );

        match result {
            Ok(jstr) => {
                let jstr_obj = match jstr.l() {
                    Ok(obj) => obj,
                    Err(_) => return None,
                };
                
                if jstr_obj.is_null() {
                    None
                } else {
                    let jstr = JString::from(jstr_obj);
                    env.get_string(&jstr)
                        .ok()
                        .map(|jstring| jstring.to_string_lossy().to_string())
                }
            }
            Err(_) => None,
        }
    }
}

/// Java AgentResponseStream 的 JNI 包装
/// 注意：这个结构持有 Java 对象的全局引用
struct JavaAgentResponseStream {
    obj: GlobalRef,
}

unsafe impl Send for JavaAgentResponseStream {}
unsafe impl Sync for JavaAgentResponseStream {}

impl JavaAgentResponseStream {
    fn new(env: &mut JNIEnv, obj: JObject) -> Result<Self, jni::errors::Error> {
        let global_ref = env.new_global_ref(obj)?;
        Ok(JavaAgentResponseStream { obj: global_ref })
    }

    fn update(&self, env: &mut JNIEnv, content: &str) -> Result<(), jni::errors::Error> {
        let jstr = env.new_string(content)?;
        let jobj: JObject = jstr.into();
        env.call_method(
            self.obj.as_obj(),
            "update",
            "(Ljava/lang/String;)V",
            &[JValue::Object(&jobj)],
        )?;
        Ok(())
    }
}

/// JNI 函数：创建 QAAgent 实例
/// 
/// # Arguments
/// * `env` - JNI 环境
/// * `_class` - Java 类对象
/// * `workspace` - workspace 字符串
/// 
/// # Returns
/// 返回指向 QAAgentWrapper 的指针（jlong）
#[no_mangle]
pub extern "system" fn Java_QAAgent_nativeNew(
    mut env: JNIEnv,
    _class: JClass,
    workspace: JString,
) -> jlong {
    let workspace_str: String = env
        .get_string(&workspace)
        .expect("Couldn't get Java string!")
        .into();

    let agent = QAAgent::new(workspace_str);
    let wrapper = QAAgentWrapper {
        agent: Arc::new(agent),
    };

    Box::into_raw(Box::new(wrapper)) as jlong
}

/// JNI 函数：调用 answer 方法
/// 
/// # Arguments
/// * `env` - JNI 环境
/// * `_class` - Java 类对象
/// * `ptr` - QAAgent 指针
/// * `query` - 用户查询字符串
/// * `model` - LLModel Java 对象
/// * `stream` - AgentResponseStream Java 对象
#[no_mangle]
pub extern "system" fn Java_QAAgent_nativeAnswer(
    mut env: JNIEnv,
    _class: JClass,
    ptr: jlong,
    query: JString,
    model: JObject,
    stream: JObject,
) {
    if ptr == 0 {
        eprintln!("Invalid QAAgent pointer");
        return;
    }
    
    let wrapper = unsafe {
        &*(ptr as *const QAAgentWrapper)
    };

    let query_str: String = match env.get_string(&query) {
        Ok(s) => s.into(),
        Err(e) => {
            eprintln!("Failed to get query string: {:?}", e);
            return;
        }
    };

    // 创建 Java 对象的包装
    let mut java_model = match JavaLLModel::new(&mut env, model) {
        Ok(m) => m,
        Err(e) => {
            eprintln!("Failed to create JavaLLModel: {:?}", e);
            return;
        }
    };

    let java_stream = match JavaAgentResponseStream::new(&mut env, stream) {
        Ok(s) => s,
        Err(e) => {
            eprintln!("Failed to create JavaAgentResponseStream: {:?}", e);
            return;
        }
    };

    // 直接在 JNI 函数中实现 answer 逻辑，避免复杂的生命周期问题
    // 首先输出 workspace 和 query
    if let Err(e) = java_stream.update(&mut env, &format!("Workspace: {}\n", wrapper.agent.workspace)) {
        eprintln!("Failed to update stream: {:?}", e);
        return;
    }
    
    if let Err(e) = java_stream.update(&mut env, &format!("Query: {}\n", query_str)) {
        eprintln!("Failed to update stream: {:?}", e);
        return;
    }
    
    if let Err(e) = java_stream.update(&mut env, "Processing...\n") {
        eprintln!("Failed to update stream: {:?}", e);
        return;
    }

    // 从模型获取流式输出
    while let Some(output) = java_model.next(&mut env) {
        if let Err(e) = java_stream.update(&mut env, &format!("Model output: {}\n", output)) {
            eprintln!("Failed to update stream: {:?}", e);
            return;
        }
    }

    if let Err(e) = java_stream.update(&mut env, "Done.\n") {
        eprintln!("Failed to update stream: {:?}", e);
    }
}

/// JNI 函数：释放 QAAgent 实例
/// 
/// # Arguments
/// * `_env` - JNI 环境
/// * `_class` - Java 类对象
/// * `ptr` - QAAgent 指针
#[no_mangle]
pub extern "system" fn Java_QAAgent_nativeFree(
    _env: JNIEnv,
    _class: JClass,
    ptr: jlong,
) {
    if ptr != 0 {
        unsafe {
            let _ = Box::from_raw(ptr as *mut QAAgentWrapper);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // Mock LLModel 实现用于测试
    struct MockLLModel {
        count: usize,
        max_count: usize,
    }

    impl LLModel for MockLLModel {
        fn next(&mut self) -> Option<String> {
            if self.count < self.max_count {
                self.count += 1;
                Some(format!("hello world {}", self.count))
            } else {
                None
            }
        }
    }

    // Mock AgentResponseStream 实现用于测试
    struct MockStream;

    impl AgentResponseStream for MockStream {
        fn update(&self, content: &str) {
            // 在测试中，我们可以收集输出
            println!("Stream update: {}", content);
        }
    }

    #[test]
    fn test_qa_agent() {
        let agent = QAAgent::new("test_workspace".to_string());
        let mut model = MockLLModel {
            count: 0,
            max_count: 5,
        };
        let stream = MockStream;

        agent.answer("test query".to_string(), &mut model, &stream);
    }
}

